resampleInStratum <- function(val, data, strata) {
  ok <- which(data$strata == val)
  sample(ok, length(ok), replace = TRUE)
}

resample <- function(data) {
  vals <- unique(data$strata)
  do.call(c, lapply(vals, resampleInStratum, data))
}

boot <- function(formula, FUN, data, fitted_model, B) {
  resamples <- do.call(cbind, lapply(1:B, function(b) resample(data)))
  boots <- matrix(NA, nrow = B, ncol = length(coef(fitted_model)))
  boots <- as.data.frame(boots)
  names(boots) <- names(coef(fitted_model))
  for (b in 1:B) {
    resampled_data <- data[resamples[, b],]
    coefs <- coef(FUN(formula, data = resampled_data, weights = weights))
    for (var in names(coefs)) {
      boots[b, var] <- coefs[var]
    }
  }
  boots
}

cace <- function(response, baseline, data, B = 5, strata = NULL, ipw = FALSE,
  conditioning_covariate = NULL, incl_sameparty = FALSE) {

  y <- data[, paste(response, '01', sep = '.')]
  data <- data[!is.na(y), ]

  if (!ipw) {
    data$weights <- 1
  } else {
    data$weights <- data[, paste('ipw', response, sep = '.')]
  }

  if (is.null(strata)) {
    data$strata <- 1
  }  else {
    data$strata <- data[, strata]
  }

  if (is.null(conditioning_covariate)) {
    if (incl_sameparty) {
      formula <- paste(response, '.01',
                   '~attended.session + same.party +', baseline, 
                   '.NArecode | . - attended.session + assigned.to.session', 
                   sep = '')
    } else {
      formula <- paste(response, '.01',
                   '~attended.session +', baseline, 
                   '.NArecode | . - attended.session + assigned.to.session', 
                   sep = '')
    }
  } else {
    formula <- paste(response, '.01',
                 '~attended.session * same.party +', baseline, 
                 '.NArecode | . - attended.session * same.party + assigned.to.session * same.party', 
                 sep = '')
  }
  fitted_model <- ivreg(as.formula(formula), data = data, weights = weights)
  boots <- boot(formula, ivreg, data, fitted_model, B)
  out <- list(treat = 'attended.session', fitted_model = fitted_model, boots = boots)
  class(out) <- 'est'
  out
}

itt <- function(response, baseline, data, B = 5, strata = NULL, ipw = FALSE, incl_sameparty = FALSE) {

  y <- data[, paste(response, '01', sep = '.')]
  data <- data[!is.na(y), ]

  if (!ipw) {
    data$weights <- 1
  } else {
    data$weights <- data[, paste('ipw', response, sep = '.')]
  }

  if (is.null(strata)) {
    data$strata <- 1
  }  else {
    data$strata <- data[, strata]
  }

  if (incl_sameparty) {
    formula <- paste(response, '.01',
                 '~ assigned.to.session + same.party +', baseline, 
                 '.NArecode', sep = '')
  } else {
    formula <- paste(response, '.01',
                 '~ assigned.to.session +', baseline, 
                 '.NArecode', sep = '')

  }
  fitted_model <- lm(as.formula(formula), data = data, weights = weights)
  boots <- boot(formula, lm, data, fitted_model, B)
  out <- list(treat = 'assigned.to.session', fitted_model = fitted_model, boots = boots)
  class(out) <- 'est'
  out
}

summary.est <- function(object) {
  tr <- object$treat
  fm <- object$fitted_model
  boots <- object$boots[, tr]
  B <- length(boots)
  N <- nrow(fm$model)
  EST <- coef(fm)[tr]
  SE <- bootSE(boots, N, B)
  P <- bootP(boots)
  q025 <- quantile(boots, .025)
  q975 <- quantile(boots, .975)
  out <- data.frame(B = EST, SE = SE, P = P, q025 = q025, q975 = q975, N = N)  
  rownames(out) <- terms(fm)[[2]]
  round(out, 3)
}

bootSE <- function(x, n, B) {
  sq <- (x - mean(x)) ^ 2
  sqrt(n / (n - 1)) * sqrt(sum(sq) / (B - 1))
}

bootP <- function(x) {
  2 * min(mean(x < 0), mean(x > 0))
}

NArecode <- function(x) {
  out <- x
  out[is.na(x)] <- -1
  factor(out)
}

ConditionalTrimmingBounds <- function(y, x, z) {
  #  Calculates conditional trimming bounds by subsetting by a covariate,
  #    performing trimming bounds analysis for values of y, z at each
  #    covariate value, and averaging results
  #
  #  Args:
  #    y: response variable
  #    x: covariate
  #    z: treatment variable
  #
  #  Returns:
  #    a pair of bounds on the treatment effect
  x.recode <- x  # Recode missing covariate values and coerce to factor
  x.recode[is.na(x)] <- max(x, na.rm=TRUE) + 1
  x.recode <- factor(x.recode)

  # Function to calc trimming bounds for (y, z) at value of x.recode
  f <- function(val) TrimmingBounds(y[x.recode == val], z[x.recode == val])

  # Apply f at each value of x.recode, bind results
  tbs <- do.call(rbind, lapply(levels(x.recode), f))

  # Drop undefined bounds, calculate weighted means, return
  w <- tabulate(x.recode) / length(x.recode)  # Calculate weights
  c(mean(tbs[, 1], weights = w, na.rm = TRUE),
    mean(tbs[, 2], weights = w, na.rm = TRUE))
}

ComplierReporter <- function(response, baseline, data) {
  data <- subset(data, assigned.to.session == 1)
  y1 <- data[, response]
  y0 <- data[, baseline]
  cr <- which(data$attended.session == 1 & !is.na(y1))
  cn <- which(data$attended.session == 1 &  is.na(y1))
  nr <- which(data$attended.session == 0 & !is.na(y1))
  nn <- which(data$attended.session == 0 &  is.na(y1))
  type <- rep('', nrow(data))
  type[cr] <- 'cr'
  type[cn] <- 'cn'
  type[nr] <- 'nr'
  type[nn] <- 'nn'
  out <- data.frame(mean.cr = rep(NA, 2), 
                    mean.cn = rep(NA, 2), 
                    mean.nr = rep(NA, 2), 
                    mean.nn = rep(NA, 2), 
                    Fstat = rep(NA, 2), 
                    Pvalue = rep(NA, 2))
  out[1, 1:4] <- c(mean(y0[cr], na.rm = TRUE), 
                   mean(y0[cn], na.rm = TRUE), 
                   mean(y0[nr], na.rm = TRUE), 
                   mean(y0[nn], na.rm = TRUE))
  out[2, 1:4] <- c(length(cr), 
                   length(cn), 
                   length(nr), 
                   length(nn))
  fm <- aov(y0 ~ type)
  out[1, 5] <- summary(fm)[[1]]$'F value'[1]
  out[1, 6] <- summary(fm)[[1]]$'Pr(>F)'[1]
  out
}

TrimmingBounds <- function(y, z) {
  #  Calculates trimming bounds by replacing missing values of y with
  #    extreme (lo & hi) observed values for opposite treatment value
  #
  #  Args:
  #    y: response variable
  #    z: treatment variable
  #
  #  Returns:
  #    a pair of bounds on the treatment effect

  # if all observations are either treated or control, trimming bounds
  #   are not defined
  if (sum(z) %in% c(0, length(z))) {
    return(c(NA, NA))
  }
  pT <- mean(!is.na(y[z == 1]))    # percent treated that are observed
  pC <- mean(!is.na(y[z == 0]))    # percent control that are observed
  zobs <- z[!is.na(y)]             # treatments for observed
  yobs <- y[!is.na(y)]             # responses for observed
  yTobs <- sort(yobs[zobs == 1])   # observed treated y's, sorted lo values 1st
  yCobs <- sort(yobs[zobs == 0])   # observed control y's, sorted lo values 1st
  if (pC > pT) {                   # if observed more controls than treated...
    q <- (pC - pT) / pC            # ratio of extra missingness among treated
    nobs <- length(yCobs)          # number of observed control values
    yClo <- tail(yCobs, -ceiling(q * nobs))      # trim lo control y-values
    yChi <- head(yCobs, ceiling((1 - q) * nobs)) # trim hi control y-values
    lb <- mean(yTobs) - mean(yClo) # average effect trimming lo control y-values
    ub <- mean(yTobs) - mean(yChi) # average effect trimming hi control y-values
  } else {                         # if observed more controls than treated...
    q <- (pT - pC) / pT            # ratio of extra missingness among control
    nobs <- length(yTobs)          # number of observed treated values
    yTlo <- tail(yTobs, -ceiling(q * nobs))      # trim lo treated y-values
    yThi <- head(yTobs, ceiling((1 - q) * nobs)) # trim hi treated y-values
    lb <- mean(yThi) - mean(yCobs) # average effect trimming hi treated y-values
    ub <- mean(yTlo) - mean(yCobs) # average effect trimming lo treated y-values
  }
  # return bounds
  if (is.na(lb) | is.na(ub)) {
    c(NA, NA)
  } else {
    if (lb < ub) {
      c(lb, ub)
    } else {
      c(ub, lb)
    }
  }
}

summarize <- function(x) {
  c(mean(x, na.rm = TRUE), sd(x, na.rm = TRUE), sum(is.na(x)))
}

scale01 <- function(x) {
  (x - min(x, na.rm = TRUE)) /  (max(x, na.rm = TRUE) - min(x, na.rm = TRUE))
}

loo <- function(response, baseline, data = d1, strata = 'congressional.district') {
  # leave one district out ivreg
  formula <- paste(response, '.01',
                '~attended.session + same.party +', baseline, 
                '.NArecode | . - attended.session + assigned.to.session', 
                sep = '')
  formula <- as.formula(formula)
  fitted_model <- ivreg(formula, data = data)
  data <- data[-fitted_model$na.action, ]
  n <- nrow(data)
  if (!is.null(strata)) {
    strata <- data[, strata]
    vals <- unique(strata)
  } else {
    strata <- 1
    vals <- 1
  }
  LOOEST <- function(val) {
    loo_data <- subset(data, strata != val)
    coef(ivreg(formula, data = loo_data))
  }
  loos <- do.call(rbind, lapply(vals, LOOEST))
  counts <- do.call(c, lapply(vals, function(val) sum(strata == val)))
  list(fitted_model = fitted_model, loos = loos, counts = counts, n = n)
}

SummarizeLOO <- function(object) {
  treat <- 'attended.session'
  LOO_ESTs <- object$loos[, treat]
  EST <- coef(object$fitted_model)[treat]
  LOO_SE <- sqrt(sum((1 - object$counts / object$n) * (EST - LOO_ESTs) ^ 2)) 
  LOO_P <- 2 * (1 - pnorm(abs(EST / LOO_SE)))
  out <- data.frame(EST = EST, SE = LOO_SE, P = LOO_P, N = object$n)
  rownames(out) <- terms(object$fitted_model)[[2]]
  round(out, 3)
}

SummarizeConditionalEffects <- function(object) {
  coefs <- coef(object$fitted_model)
  N <- nrow(object$fitted_model$model)
  EST1 <- coefs['attended.session'] + coefs['attended.session:same.party']
  EST2 <- coefs['attended.session']
  boots1 <- object$boots[, 'attended.session'] + object$boots[, 'attended.session:same.party']
  boots2 <- object$boots[, 'attended.session']
  B <- length(boots1)
  SE1 <- bootSE(boots1, N, B)
  SE2 <- bootSE(boots2, N, B)
  P1 <- bootP(boots1)
  P2 <- bootP(boots2)
  q025 <- c(quantile(boots1, .025), quantile(boots2, .025))
  q975 <- c(quantile(boots1, .975), quantile(boots2, .975))
  P_compare <- 2 * min(mean(boots2 < boots1), mean(boots2 > boots1))
  results <- data.frame(EST = c(EST1, EST2), SE = c(SE1, SE2), P = c(P1, P2),
    q025 = q025, q975 = q975, N = c(N, NA), P_compare = c(P_compare, NA))
  round(results, 3)
}